#include "viml/compiler.h"
#include "viml/ast.h"
#include "viml/scope_manager.h"
#include "viml/simple_optimizer.h"
#include "viml/optimization_pipeline.h"
#include "viml/optimization_passes.h"
#include <stdexcept>
#include <cmath>
#include <iostream>

namespace viml {

Compiler::Compiler() 
    : optimizer_(std::make_unique<SimpleOptimizer>()),
      pipeline_(std::make_unique<OptimizationPipeline>())
{
    reset();
    
    // 初始化优化管道(可选,不影响现有功能)
    // 如果想使用Pipeline模式,可以通过get_optimization_pipeline()配置
}

void Compiler::reset() {
    instructions.clear();
    constants.clear();
    names.clear();
    instruction_lines.clear();
    current_line = 0;
    // 注意：不要清空functions列表，因为函数对象需要在编译后仍然可用
    // functions.clear();
    context = CompileContext();
}

uint8_t Compiler::allocate_register() {
    if (context.next_register >= 255) {
        throw std::runtime_error("Too many registers needed");
    }
    return context.next_register++;
}

uint16_t Compiler::add_constant(const Value& value) {
    // 检查是否已存在相同常量
    for (size_t i = 0; i < constants.size(); ++i) {
        if (constants[i] == value) {
            return static_cast<uint16_t>(i);
        }
    }
    
    constants.push_back(value);
    return static_cast<uint16_t>(constants.size() - 1);
}

uint16_t Compiler::add_name(const std::string& name) {
    names.push_back(name);
    return names.size() - 1;
}

void Compiler::emit(OpCode op, uint8_t a, uint8_t b, uint8_t c) {
    instructions.emplace_back(op, a, b, c);
    instruction_lines.push_back(static_cast<int>(current_line));
}

void Compiler::emit_abx(OpCode op, uint8_t a, uint16_t bx) {
    instructions.emplace_back(op, a, bx);
    instruction_lines.push_back(static_cast<int>(current_line));
}

uint8_t Compiler::get_or_create_local(const std::string& name) {
    auto it = context.local_variables.find(name);
    if (it != context.local_variables.end()) {
        return it->second;
    }
    
    uint8_t reg = allocate_register();
    context.local_variables[name] = reg;
    return reg;
}

void Compiler::emit_scope_load(uint8_t reg, ScopeId scope, uint16_t name_idx) {
    uint8_t scope_id = static_cast<uint8_t>(scope);
    uint16_t bx = (static_cast<uint16_t>(scope_id & 0xF) << 12) | (name_idx & 0xFFF);
    emit_abx(OpCode::LOAD_SCOPE, reg, bx);
}

void Compiler::emit_scope_store(ScopeId scope, uint16_t name_idx, uint8_t value_reg) {
    uint8_t scope_id = static_cast<uint8_t>(scope);
    uint16_t bx = (static_cast<uint16_t>(scope_id & 0xF) << 12) | (name_idx & 0xFFF);
    emit_abx(OpCode::STORE_SCOPE, value_reg, bx);
}

std::pair<ScopeId, std::string> Compiler::parse_variable_scope(const std::string& name) {
    // 解析作用域前缀
    if (name.length() >= 3 && name[1] == ':') {
        ScopeId scope = parse_scope_prefix(name[0]);
        std::string var_name = name.substr(2);
        return {scope, var_name};
    }
    
    // 没有前缀，根据上下文决定作用域
    if (context.in_function) {
        // 在函数内部，无前缀变量默认为局部作用域
        // 但为了兼容性，我们使用全局作用域（简化处理）
        return {ScopeId::GLOBAL, name};
    }
    
    // 顶层默认是全局作用域
    return {ScopeId::GLOBAL, name};
}

std::shared_ptr<Bytecode> Compiler::compile(const std::string& source, const std::string& filename) {
    reset();
    
    try {
        // 词法分析
        Lexer lexer(source);
        auto tokens = lexer.tokenize();
        
        // 语法分析
        Parser parser(tokens);
        auto program = parser.parse();

        // AST优化（启用）
        std::shared_ptr<Program> program_ptr;
        if (optimizer_ && optimizer_->is_enabled()) {
            // 将 unique_ptr 转换为 shared_ptr 进行优化
            program_ptr = std::shared_ptr<Program>(program.release());
            program_ptr = optimizer_->optimize_constants_in_ast(program_ptr);
            program_ptr = optimizer_->eliminate_dead_code(program_ptr);
        } else {
            program_ptr = std::shared_ptr<Program>(program.release());
        }

        // 代码生成
        program_ptr->accept(*this);
        
        // 添加HALT指令
        emit(OpCode::HALT);
        
        // 创建字节码对象
        auto bytecode = std::make_shared<Bytecode>(filename);

        // 设置常量池、名字和指令
        for (const auto& constant : constants) {
            bytecode->add_constant(constant);
        }

        for (const auto& name : names) {
            bytecode->add_name(name);
        }

        for (size_t i = 0; i < instructions.size(); ++i) {
            bytecode->add_instruction(instructions[i].opcode(), instructions[i].a(), instructions[i].bx());
            if (i < instruction_lines.size()) {
                bytecode->add_line_number(instruction_lines[i]);
            } else {
                bytecode->add_line_number(1); // 回退到简化行号映射
            }
        }

        // 字节码优化（启用）
        if (optimizer_ && optimizer_->is_enabled()) {
            optimizer_->optimize_registers_in_bytecode(bytecode);
        }
        
        // 调试输出
        // std::cout << "Generated bytecode for " << filename << ":\n";
        // bytecode->dump();
        
        return bytecode;
        
    } catch (const std::exception& e) {
        throw std::runtime_error("Compilation failed: " + std::string(e.what()));
    }
}

// AST访问者实现

void Compiler::visit(NumberLiteral& node) {
    current_line = node.line;

    // 根据VimL类型系统区分整数和浮点数
    Value value;
    if (std::floor(node.value) == node.value &&
        node.value >= -2147483648.0 && node.value <= 2147483647.0) {
        // 整数值，创建 NUMBER 类型
        value = Value::make_number(static_cast<int64_t>(node.value));
    } else {
        // 浮点数值，创建 FLOAT 类型
        value = Value::make_float(node.value);
    }

    uint16_t const_idx = add_constant(value);
    uint8_t reg = allocate_register();
    emit_abx(OpCode::LOAD_CONST, reg, const_idx);
}

void Compiler::visit(StringLiteral& node) {
    current_line = node.line;

    uint16_t const_idx = add_constant(Value::string(node.value));
    uint8_t reg = allocate_register();
    emit_abx(OpCode::LOAD_CONST, reg, const_idx);
}

void Compiler::visit(Identifier& node) {
    current_line = node.line;

    // 检查是否是环境变量（以$开头）
    if (node.name.length() >= 2 && node.name[0] == '$') {
        // 环境变量，使用ENV作用域
        std::string var_name = node.name.substr(1);  // 去掉$前缀
        
        uint16_t name_idx = 0;
        bool found = false;
        for (size_t i = 0; i < names.size(); ++i) {
            if (names[i] == var_name) {
                name_idx = static_cast<uint16_t>(i);
                found = true;
                break;
            }
        }
        if (!found) {
            name_idx = add_name(var_name);
        }
        
        uint8_t reg = allocate_register();
        emit_scope_load(reg, ScopeId::ENV, name_idx);
    }
    // 检查是否有显式作用域前缀
    else if (node.name.length() >= 3 && node.name[1] == ':') {
        // 有显式作用域前缀，使用新的作用域系统
        auto [scope, var_name] = parse_variable_scope(node.name);
        
        uint16_t name_idx = 0;
        bool found = false;
        for (size_t i = 0; i < names.size(); ++i) {
            if (names[i] == var_name) {
                name_idx = static_cast<uint16_t>(i);
                found = true;
                break;
            }
        }
        if (!found) {
            name_idx = add_name(var_name);
        }
        
        uint8_t reg = allocate_register();
        emit_scope_load(reg, scope, name_idx);
    } else {
        // 检查是否是局部变量（包括函数参数和for循环变量）
        auto local_it = context.local_variables.find(node.name);
        if (local_it != context.local_variables.end()) {
            // 是局部变量或参数，直接从寄存器加载
            uint8_t reg = allocate_register();
            emit(OpCode::LOAD_LOCAL, reg, local_it->second);
            return;
        }
        
        // 检查是否在函数内部
        if (context.in_function) {
            // 在函数内部但不是局部变量，回退到全局变量系统
        }
        
        // 使用旧的全局变量系统（兼容模式）
        uint16_t name_idx = 0;
        bool found = false;
        for (size_t i = 0; i < names.size(); ++i) {
            if (names[i] == node.name) {
                name_idx = static_cast<uint16_t>(i);
                found = true;
                break;
            }
        }
        if (!found) {
            name_idx = add_name(node.name);
        }
        uint8_t reg = allocate_register();
        emit_abx(OpCode::LOAD_GLOBAL, reg, name_idx);
    }
}

void Compiler::visit(SpecialConstantNode& node) {
    Value value;
    
    // 根据特殊常量名称创建相应的值
    if (node.constant_name == "v:true") {
        value = Value::boolean(true);
    } else if (node.constant_name == "v:false") {
        value = Value::boolean(false);
    } else if (node.constant_name == "v:null") {
        value = Value::nil();
    } else if (node.constant_name == "v:none") {
        value = Value::none();
    } else if (node.constant_name == "v:errmsg") {
        value = Value::string(""); // 默认空错误信息
    } else if (node.constant_name == "v:t_number") {
        value = Value::make_number(0);
    } else if (node.constant_name == "v:t_string") {
        value = Value::make_number(1);
    } else if (node.constant_name == "v:t_function") {
        value = Value::make_number(2);
    } else if (node.constant_name == "v:t_list") {
        value = Value::make_number(3);
    } else if (node.constant_name == "v:t_dict") {
        value = Value::make_number(4);
    } else if (node.constant_name == "v:t_float") {
        value = Value::make_number(5);
    } else if (node.constant_name == "v:t_bool") {
        value = Value::make_number(6);
    } else {
        throw std::runtime_error("Unknown special constant: " + node.constant_name);
    }
    
    uint16_t const_idx = add_constant(value);
    uint8_t reg = allocate_register();
    emit_abx(OpCode::LOAD_CONST, reg, const_idx);
}

void Compiler::visit(BinaryExpression& node) {
    // 计算左操作数
    node.left->accept(*this);
    uint8_t left_reg = context.next_register - 1;
    
    // 计算右操作数
    node.right->accept(*this);
    uint8_t right_reg = context.next_register - 1;
    
    // 生成结果寄存器
    uint8_t result_reg = allocate_register();
    
    // 生成操作指令
    switch (node.op) {
        case BinaryExpression::Operator::ADD:
            emit(OpCode::ADD, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::SUB:
            emit(OpCode::SUB, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::MUL:
            emit(OpCode::MUL, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::DIV:
            emit(OpCode::DIV, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::MOD:
            emit(OpCode::MOD, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::EQ:
            emit(OpCode::EQ, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::NE:
            emit(OpCode::NE, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::LT:
            emit(OpCode::LT, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::LE:
            emit(OpCode::LE, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::GT:
            emit(OpCode::GT, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::GE:
            emit(OpCode::GE, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::AND:
            emit(OpCode::AND, result_reg, left_reg, right_reg);
            break;
        case BinaryExpression::Operator::OR:
            emit(OpCode::OR, result_reg, left_reg, right_reg);
            break;
        default:
            throw std::runtime_error("Unsupported binary operator");
    }
}

void Compiler::visit(UnaryExpression& node) {
    // 计算操作数
    node.operand->accept(*this);
    uint8_t operand_reg = context.next_register - 1;
    
    // 生成结果寄存器
    uint8_t result_reg = allocate_register();
    
    // 生成操作指令
    switch (node.op) {
        case UnaryExpression::Operator::MINUS:
            emit(OpCode::NEG, result_reg, operand_reg);
            break;
        case UnaryExpression::Operator::NOT:
            emit(OpCode::NOT, result_reg, operand_reg);
            break;
        default:
            throw std::runtime_error("Unsupported unary operator");
    }
}

void Compiler::visit(CallExpression& node) {
    // 编译函数表达式（通常是标识符）
    node.callee->accept(*this);
    uint8_t func_reg = context.next_register - 1;
    
    // 保存参数开始的寄存器索引
    uint8_t param_start_reg = context.next_register;
    
    // 编译参数并在编译后立即将其移动到正确位置
    for (size_t i = 0; i < node.arguments.size(); ++i) {
        // 编译参数表达式
        node.arguments[i]->accept(*this);
        uint8_t arg_result_reg = context.next_register - 1;
        
        // 计算参数应该放置的目标寄存器
        uint8_t target_reg = func_reg + 1 + static_cast<uint8_t>(i);
        
        // 如果结果不在正确位置，移动它
        if (arg_result_reg != target_reg) {
            emit(OpCode::MOVE, target_reg, arg_result_reg);
        }
    }
    uint8_t argc = static_cast<uint8_t>(node.arguments.size());
    
    // 生成结果寄存器
    uint8_t result_reg = allocate_register();
    
    // 生成CALL指令：CALL result_reg, func_reg, argc
    // 参数在func_reg+1..func_reg+argc
    emit(OpCode::CALL, result_reg, func_reg, argc);
    
    // 恢复寄存器状态，确保返回的是结果寄存器
    context.next_register = result_reg + 1;
}

void Compiler::visit(ListLiteral& node) {
    // 保存当前寄存器状态
    uint8_t original_next_reg = context.next_register;
    
    // 分配列表寄存器
    uint8_t list_reg = allocate_register();
    
    // 创建空列表
    uint8_t capacity = static_cast<uint8_t>(std::min(node.elements.size(), static_cast<size_t>(255)));
    emit(OpCode::NEW_LIST, list_reg, capacity);
    
    // 添加列表元素
    for (auto& element : node.elements) {
        // 计算元素表达式
        element->accept(*this);
        uint8_t element_reg = context.next_register - 1;
        
        // 将元素添加到列表
        emit(OpCode::LIST_APPEND, list_reg, element_reg);
    }
    
    // 恢复寄存器状态，确保返回的是列表寄存器
    context.next_register = original_next_reg + 1;  // list_reg是original_next_reg
}

void Compiler::visit(DictLiteral& node) {
    // 保存当前寄存器状态
    uint8_t original_next_reg = context.next_register;
    
    // 分配字典寄存器
    uint8_t dict_reg = allocate_register();
    
    // 创建空字典
    emit(OpCode::NEW_DICT, dict_reg);
    
    // 添加字典键值对
    for (const auto& pair : node.pairs) {
        // 计算键表达式（通常是字符串字面量）
        pair.first->accept(*this);
        uint8_t key_reg = context.next_register - 1;
        
        // 计算值表达式
        pair.second->accept(*this);
        uint8_t value_reg = context.next_register - 1;
        
        // 将键值对添加到字典
        emit(OpCode::DICT_SET, dict_reg, key_reg, value_reg);
    }
    
    // 恢复寄存器状态，确保返回的是字典寄存器
    context.next_register = original_next_reg + 1;  // dict_reg是original_next_reg
}

void Compiler::visit(IndexExpression& node) {
    // 编译对象表达式
    node.object->accept(*this);
    uint8_t obj_reg = context.next_register - 1;
    
    // 编译索引表达式
    node.index->accept(*this);
    uint8_t index_reg = context.next_register - 1;
    
    // 生成结果寄存器
    uint8_t result_reg = allocate_register();
    
    // 生成GET_INDEX指令
    emit(OpCode::GET_INDEX, result_reg, obj_reg, index_reg);
}

void Compiler::visit(LetStatement& node) {
    // 计算赋值表达式
    node.value->accept(*this);
    uint8_t value_reg = context.next_register - 1;
    
    // 检查是否有显式作用域前缀
    if (node.name.length() >= 3 && node.name[1] == ':') {
        // 有显式作用域前缀，使用新的作用域系统
        auto [scope, var_name] = parse_variable_scope(node.name);
        
        uint16_t name_idx = 0;
        bool found = false;
        for (size_t i = 0; i < names.size(); ++i) {
            if (names[i] == var_name) {
                name_idx = static_cast<uint16_t>(i);
                found = true;
                break;
            }
        }
        if (!found) {
            name_idx = add_name(var_name);
        }
        
        emit_scope_store(scope, name_idx, value_reg);
    } else {
        // 没有显式前缀，检查是否在函数内部
        uint16_t name_idx = 0;
        bool found = false;
        for (size_t i = 0; i < names.size(); ++i) {
            if (names[i] == node.name) {
                name_idx = static_cast<uint16_t>(i);
                found = true;
                break;
            }
        }
        if (!found) {
            name_idx = add_name(node.name);
        }
        
        if (context.in_function) {
            // 在函数内部，将变量存储到局部寄存器
            uint8_t var_reg = get_or_create_local(node.name);
            
            if (var_reg != value_reg) {
                emit(OpCode::MOVE, var_reg, value_reg);
            }
            
            // 同时存储到局部作用域
            emit_scope_store(ScopeId::LOCAL, name_idx, var_reg);
        } else {
            // 不在函数内部，使用旧的全局变量系统（兼容模式）
            uint8_t var_reg = get_or_create_local(node.name);
            
            if (var_reg != value_reg) {
                emit(OpCode::MOVE, var_reg, value_reg);
            }
            
            emit_abx(OpCode::STORE_GLOBAL, var_reg, name_idx);
        }
    }
}

void Compiler::visit(EchoStatement& node) {
    current_line = node.line;

    // 检查表达式是否为列表字面量，如果是，则分别处理每个元素
    if (auto list_literal = dynamic_cast<ListLiteral*>(node.expression.get())) {
        // 处理多个表达式的情况
        uint8_t start_reg = context.next_register;
        
        // 为每个元素分配寄存器并编译
        for (size_t i = 0; i < list_literal->elements.size(); ++i) {
            // 编译元素表达式
            list_literal->elements[i]->accept(*this);
            uint8_t element_result_reg = context.next_register - 1;
            
            // 将结果移动到连续的寄存器位置
            if (element_result_reg != start_reg + i) {
                emit(OpCode::MOVE, start_reg + i, element_result_reg);
            }
            
            // 重置寄存器状态，确保下一个元素从正确位置开始
            context.next_register = start_reg + i + 1;
        }
        
        // 计算元素数量
        uint8_t count = static_cast<uint8_t>(list_literal->elements.size());
        
        // 生成ECHO指令，输出所有元素，空格分隔但不换行到最后
        emit(OpCode::ECHO, start_reg, count);
        
        // 确保程序返回NIL值
        emit(OpCode::LOAD_NIL, 0);
        
        // 恢复寄存器状态
        context.next_register = start_reg + 1;
    } else {
        // 处理单个表达式的情况
        node.expression->accept(*this);
        uint8_t expr_reg = context.next_register - 1;
        emit(OpCode::ECHO, expr_reg, 1);
        // 确保程序返回NIL值
        emit(OpCode::LOAD_NIL, 0);
        // 确保寄存器状态正确
        context.next_register = expr_reg + 1;
    }
}

void Compiler::visit(IfStatement& node) {
    // 计算条件表达式
    node.condition->accept(*this);
    uint8_t condition_reg = context.next_register - 1;
    
    // 生成JMP_IF_FALSE指令：如果条件为假，跳转到else部分或endif
    size_t jmp_if_false_addr = instructions.size();
    emit_abx(OpCode::JMP_IF_FALSE, condition_reg, 0); // 占位符，稍后填充
    
    // 编译then部分
    for (auto& stmt : node.then_body) {
        stmt->accept(*this);
    }
    
    size_t jmp_end_addr = 0;
    bool has_else = !node.else_body.empty();
    
    if (has_else) {
        // 如果有else部分，生成JMP指令跳过else部分
        jmp_end_addr = instructions.size();
        emit_abx(OpCode::JMP, 0, 0); // 占位符，稍后填充
    }
    
    // 计算JMP_IF_FALSE的目标地址（else部分开始或endif）
    size_t else_start_addr = instructions.size();
    
    if (has_else) {
        // 编译else部分
        for (auto& stmt : node.else_body) {
            stmt->accept(*this);
        }
    }
    
    // 计算结束地址
    size_t end_addr = instructions.size();
    
    // 回填JMP_IF_FALSE的跳转偏移
    // 偏移计算：目标地址 - 当前指令地址
    int16_t false_offset = static_cast<int16_t>(else_start_addr - jmp_if_false_addr);
    uint16_t false_bx = static_cast<uint16_t>(false_offset + 32767);
    instructions[jmp_if_false_addr] = Instruction(OpCode::JMP_IF_FALSE, condition_reg, false_bx);
    
    // 如果有else部分，回填JMP的跳转偏移
    if (has_else) {
        int16_t end_offset = static_cast<int16_t>(end_addr - jmp_end_addr);
        uint16_t end_bx = static_cast<uint16_t>(end_offset + 32767);
        instructions[jmp_end_addr] = Instruction(OpCode::JMP, 0, end_bx);
    }
}

void Compiler::visit(WhileStatement& node) {
    // 记录循环开始位置（条件检查点）
    size_t loop_start = instructions.size();
    
    // 编译条件表达式
    node.condition->accept(*this);
    uint8_t condition_reg = context.next_register - 1;
    
    // 生成条件跳转指令占位符（如果条件为假则跳转到循环结束）
    size_t jmp_if_false_addr = instructions.size();
    emit_abx(OpCode::JMP_IF_FALSE, condition_reg, 0); // 占位符
    
    // 编译循环体
    for (auto& stmt : node.body) {
        stmt->accept(*this);
    }
    
    // 生成无条件跳转指令，跳转回循环开始
    // 跳转偏移量计算：目标地址 - 当前地址
    // 根据设计文档，虚拟机执行跳转时已正确处理PC递增逻辑，因此不需要额外减1
    int16_t back_jump_offset = static_cast<int16_t>(static_cast<int64_t>(loop_start) - static_cast<int64_t>(instructions.size()));
    emit_abx(OpCode::JMP, 0, static_cast<uint16_t>(back_jump_offset + 32767));
    
    // 计算循环结束位置
    size_t loop_end = instructions.size();
    
    // 回填条件跳转指令的偏移量
    // 跳转偏移量计算：目标地址 - 当前地址
    // 根据设计文档，虚拟机执行跳转时已正确处理PC递增逻辑，因此不需要额外减1
    int16_t false_offset = static_cast<int16_t>(static_cast<int64_t>(loop_end) - static_cast<int64_t>(jmp_if_false_addr));
    uint16_t false_bx = static_cast<uint16_t>(false_offset + 32767);
    instructions[jmp_if_false_addr] = Instruction(OpCode::JMP_IF_FALSE, condition_reg, false_bx);
}

void Compiler::visit(ForStatement& node) {
    // 编译迭代对象表达式
    uint8_t iterable_reg;
    
    // 如果迭代对象是变量，直接使用其寄存器
    if (auto identifier = dynamic_cast<Identifier*>(node.iterable.get())) {
        // 首先尝试从局部变量查找
        auto it = context.local_variables.find(identifier->name);
        if (it != context.local_variables.end()) {
            iterable_reg = it->second;
        } else {
            // 如果局部变量中找不到，尝试从全局变量加载
            uint8_t global_reg = allocate_register();
            uint16_t name_idx = add_name(identifier->name);
            emit_abx(OpCode::LOAD_GLOBAL, global_reg, name_idx);
            iterable_reg = global_reg;
        }
    } else {
        // 对于复杂表达式，正常编译
        node.iterable->accept(*this);
        iterable_reg = context.next_register - 1;
    }
    
    // 分配迭代器寄存器（需要3个连续寄存器）
    uint8_t iter_reg = context.next_register;
    context.next_register += 3;
    
    // 初始化迭代器
    emit(OpCode::ITER_INIT, iter_reg, iterable_reg);
    
    // 循环开始标签（ITER_NEXT检查位置）
    size_t loop_check = instructions.size();
    
    // 生成ITER_NEXT指令，如果没有更多元素则跳转到循环结束
    size_t iter_next_addr = instructions.size();
    emit_abx(OpCode::ITER_NEXT, iter_reg, 0); // 占位符，稍后填充跳转到循环结束的偏移
    
    // 将当前元素存储到循环变量
    uint8_t var_reg = get_or_create_local(node.variable);
    emit(OpCode::MOVE, var_reg, iter_reg + 2);  // iter_reg+2 存储当前元素
    
    // 保存当前上下文
    auto saved_context = context;
    
    // 添加循环变量到局部变量表中，以便在循环体中可以正确识别
    context.local_variables[node.variable] = var_reg;
    
    // 编译循环体
    for (auto& stmt : node.body) {
        stmt->accept(*this);
    }
    
    // 恢复上下文
    context = saved_context;
    
    // 跳转回循环检查
    int16_t back_jump = static_cast<int16_t>(loop_check - instructions.size());
    emit_abx(OpCode::JMP, 0, static_cast<uint16_t>(back_jump + 32767));
    
    // 循环结束位置
    size_t loop_end = instructions.size();
    
    // 回填 ITER_NEXT 的跳转偏移（跳转到循环结束）
    int16_t end_offset = static_cast<int16_t>(loop_end - iter_next_addr);
    uint16_t end_bx = static_cast<uint16_t>(end_offset + 32767);
    instructions[iter_next_addr] = Instruction(OpCode::ITER_NEXT, iter_reg, end_bx);
}

void Compiler::visit(FunctionStatement& node) {
    // 保存当前编译状态
    auto saved_instructions = instructions;
    auto saved_constants = std::move(constants);
    auto saved_names = std::move(names);
    auto saved_context = context;
    
    // 清空编译状态以编译函数体
    instructions.clear();
    
    // 创建函数字节码对象
    auto func_bytecode = std::make_shared<Bytecode>("<func:" + node.name + ">");
    
    // 重置编译上下文以编译函数体
    instructions.clear();
    constants.clear();
    names.clear();
    context = CompileContext();
    context.in_function = true; // 设置在函数内部
    
    // 将全局名字添加到函数体的名字池中（这样函数体可以访问全局变量）
    for (const auto& global_name : saved_names) {
        add_name(global_name);
    }
    
    // 为函数参数分配寄存器（参数在寄存器0, 1, 2, ...中）
    for (size_t i = 0; i < node.parameters.size(); ++i) {
        const auto& param = node.parameters[i];
        get_or_create_local(param);
    }
    
    // 编译函数体
    for (size_t i = 0; i < node.body.size(); ++i) {
        node.body[i]->accept(*this);
    }
    
    // 如果函数未显式返回，添加默认返回 nil
    if (instructions.empty() || instructions.back().opcode() != OpCode::RETURN_OP) {
        emit(OpCode::LOAD_NIL, 0);  // 在R0中加载nil
        emit(OpCode::RETURN_OP, 0, 1);  // 返回R0
    }
    
    // 将函数字节码保存到字节码对象
    for (const auto& constant : constants) {
        func_bytecode->add_constant(constant);
    }
    for (const auto& name : names) {
        func_bytecode->add_name(name);
    }
    // 正确地将指令添加到函数字节码中
    for (const auto& instruction : instructions) {
        func_bytecode->add_instruction(instruction);
        func_bytecode->add_line_number(1);
    }
    
    // 创建函数对象
    auto function = std::make_shared<Function>(node.name, node.parameters, func_bytecode, 0);
    
    // 将函数添加到函数表（在恢复上下文之前）
    functions.push_back(function);
    uint8_t func_idx = static_cast<uint8_t>(functions.size() - 1);
    
    // 恢复编译上下文
    instructions = std::move(saved_instructions);
    constants = std::move(saved_constants);
    names = std::move(saved_names);
    context = saved_context;
    
    // 确保函数名在主程序的名字池中
    uint16_t name_idx = 0;
    bool name_found = false;
    for (size_t i = 0; i < names.size(); ++i) {
        if (names[i] == node.name) {
            name_idx = static_cast<uint16_t>(i);
            name_found = true;
            break;
        }
    }
    if (!name_found) {
        name_idx = add_name(node.name);
    }
    
    // 生成函数定义指令（跳过函数体）
    // 只跳过FUNC_DEF指令本身，让后面的LOAD_CONST和STORE_GLOBAL指令能够执行
    emit_abx(OpCode::FUNC_DEF, func_idx, 1);  // 跳过1条指令（FUNC_DEF本身）
    
    // 将函数作为值存储到全局变量（在恢复上下文之后添加常量）
    uint8_t func_reg = allocate_register();
    uint16_t func_const_idx = add_constant(Value::function(function));
    emit_abx(OpCode::LOAD_CONST, func_reg, func_const_idx);
    
    // 将函数存储到全局变量表
    emit_abx(OpCode::STORE_GLOBAL, func_reg, name_idx);
}

void Compiler::visit(ReturnStatement& node) {
    if (node.value) {
        // 编译返回值
        node.value->accept(*this);
        uint8_t return_reg = context.next_register - 1;
        emit(OpCode::RETURN_OP, return_reg, 1);
    } else {
        // 返回 nil
        uint8_t nil_reg = allocate_register();
        emit(OpCode::LOAD_NIL, nil_reg);
        emit(OpCode::RETURN_OP, nil_reg, 1);
    }
}

void Compiler::visit(CallStatement& node) {
    // 编译函数调用表达式
    node.call_expr->accept(*this);
    // 对于CallStatement，我们不需要保存返回值，但需要确保程序状态正确
    // 将NIL加载到R0中，保证程序返回NIL
    emit(OpCode::LOAD_NIL, 0);
}

void Compiler::visit(Program& node) {
    // 编译所有语句
    for (auto& stmt : node.statements) {
        stmt->accept(*this);
    }
}

// 优化器控制方法
void Compiler::enable_optimizations() {
    if (optimizer_) {
        optimizer_->set_enabled(true);
    }
}

void Compiler::disable_optimizations() {
    if (optimizer_) {
        optimizer_->set_enabled(false);
    }
}

SimpleOptimizer::OptimizationSummary Compiler::get_optimization_summary() const {
    if (optimizer_) {
        return optimizer_->get_optimization_summary();
    }
    return SimpleOptimizer::OptimizationSummary{};
}

} // namespace viml